package com.zzh.partnersys.ai.service.impl;

import com.zzh.partnersys.ai.service.DocumentParseService;
import com.zzh.partnersys.common.exception.BusinessException;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import java.io.InputStream;
import java.net.URL;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

/**
 * 文档解析服务实现
 *
 * @author zzh
 */
@Slf4j
@Service
public class DocumentParseServiceImpl implements DocumentParseService {

    @Resource
    private VectorStore elasticSearchVectorStore;

    @Override
    public void parseAndStoreDocument(Long assistantId, Long userId, String docUrl, String docName, String docType,
                                      Consumer<ParseProgress> progressCallback) {
        try {
            // 1. 根据文档类型解析文档
            List<Document> documents = parseDocument(docUrl, docName, docType, progressCallback);
            
            if (documents == null || documents.isEmpty()) {
                progressCallback.accept(new ParseProgress(0, 0, docName, "文档解析失败：无法解析文档内容", 0));
                return;
            }

            // 2. 为每个文档添加元数据（助手ID、文档名称等）
            List<Document> documentsWithMetadata = new ArrayList<>();
            String timeFormat = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));
            for (int i = 0; i < documents.size(); i++) {
                Document doc = documents.get(i);
                Map<String, Object> metadata = new HashMap<>(doc.getMetadata());
                metadata.put("assistantId", assistantId);
                metadata.put("userId", userId);
                metadata.put("docName", docName);
                metadata.put("docUrl", docUrl);
                metadata.put("docType", docType);
                metadata.put("chunkIndex", i);
                metadata.put("totalChunks", documents.size());
                metadata.put("createTime", timeFormat); // 设置创建时间
                
                // Spring AI Document 使用 getText() 而不是 getContent()
                Document docWithMetadata = new Document(doc.getText(), metadata);
                documentsWithMetadata.add(docWithMetadata);
            }

            // 3. 存储到 ES VectorStore
            progressCallback.accept(new ParseProgress(0, 0, docName, "正在存储到向量数据库...", 80));
            elasticSearchVectorStore.add(documentsWithMetadata);
            
            progressCallback.accept(new ParseProgress(0, 0, docName, "文档解析完成", 100));
            log.info("文档解析并存储成功，assistantId: {}, docName: {}, chunks: {}", assistantId, docName, documents.size());

        } catch (Exception e) {
            log.error("文档解析失败，assistantId: {}, docUrl: {}", assistantId, docUrl, e);
            progressCallback.accept(new ParseProgress(0, 0, docName, "文档解析失败：" + e.getMessage(), 0));
            throw new BusinessException("文档解析失败：" + e.getMessage(), 40000);
        }
    }

    /**
     * 根据文档类型解析文档
     */
    private List<Document> parseDocument(String docUrl, String docName, String docType,
                                         Consumer<ParseProgress> progressCallback) throws Exception {
        if (!StringUtils.hasText(docType)) {
            // 根据 URL 后缀判断类型
            docType = getFileTypeFromUrl(docUrl);
        }

        progressCallback.accept(new ParseProgress(0, 0, docName, "开始解析文档...", 10));

        String lowerType = docType.toLowerCase();
        if (lowerType.contains("pdf")) {
            return parsePdf(docUrl, progressCallback);
        } else if (lowerType.contains("md") || lowerType.contains("markdown")) {
            return parseMarkdown(docUrl, progressCallback);
        } else if (lowerType.contains("docx") || lowerType.contains("doc")) {
            return parseWord(docUrl, progressCallback);
        } else {
            throw new UnsupportedOperationException("不支持的文档类型: " + docType);
        }
    }

    /**
     * 解析 PDF 文档
     */
    private List<Document> parsePdf(String docUrl, Consumer<ParseProgress> progressCallback) {
        try {
            progressCallback.accept(new ParseProgress(0, 0, "", "正在解析 PDF 文档...", 30));
            // 使用段落模式解析（更细粒度）
//            ParagraphPdfDocumentReader pdfReader = new ParagraphPdfDocumentReader(
//                    docUrl,
//                    PdfDocumentReaderConfig.builder()
//                            .withPageTopMargin(0)
//                            .withPageExtractedTextFormatter(
//                                    ExtractedTextFormatter.builder()
//                                            .withNumberOfTopTextLinesToDelete(0)
//                                            .build()
//                            )
//                            .build()
//            );
//            List<Document> documents = pdfReader.read();
            PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(
                    docUrl, // 1. 资源路径（与 Paragraph 版本完全一致）
                    PdfDocumentReaderConfig.builder()
                            // 过滤页眉（用法相同）
                            .withPageTopMargin(0)
                            // 文本格式化（用法相同）
                            .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                    .withNumberOfTopTextLinesToDelete(0)
                                    .build())
                            //核心配置：每个 Document 对应 1 页（默认值）
                            .withPagesPerDocument(1)
                            .build());
            List<Document> rawDocuments = pdfReader.read(); // 执行解析，返回「页面级」Document 列表
            List<Document> documents = new ArrayList<>(rawDocuments.size());
            for (Document doc : rawDocuments) {
                String cleanedText = cleanPdfText(doc.getText());
                documents.add(new Document(cleanedText, doc.getMetadata()));
            }
            progressCallback.accept(new ParseProgress(0, 0, "", "PDF 解析完成", 50));
            return documents;
        } catch (Exception e) {
            log.error("PDF 解析失败，docUrl: {}", docUrl, e);
            throw new BusinessException("PDF 解析失败: " + e.getMessage(), 40000);
        }
    }

    /**
     * 清理 PDF 文本中的多余空格和换行
     */
    private String cleanPdfText(String text) {
        if (!StringUtils.hasText(text)) {
            return "";
        }
        // 统一换行符
        String normalized = text.replace("\r\n", "\n").replace("\r", "\n");
        // 将多个空格压缩为一个空格
        normalized = normalized.replaceAll("[ \\t]+", " ");
        // 去除行首行尾空格
        normalized = normalized.replaceAll("(?m)^\\s+|\\s+$", "");
        // 去除连续空行
        normalized = normalized.replaceAll("(?m)\\n{2,}", "\n\n");
        return normalized.trim();
    }

    /**
     * 解析 Markdown 文档
     */
    private List<Document> parseMarkdown(String docUrl, Consumer<ParseProgress> progressCallback) {
        try {
            progressCallback.accept(new ParseProgress(0, 0, "", "正在解析 Markdown 文档...", 30));
            MarkdownDocumentReader markdownReader = new MarkdownDocumentReader(docUrl);
            List<Document> documents = markdownReader.get();
            progressCallback.accept(new ParseProgress(0, 0, "", "Markdown 解析完成", 50));
            return documents;
        } catch (Exception e) {
            log.error("Markdown 解析失败，docUrl: {}", docUrl, e);
            throw new RuntimeException("Markdown 解析失败: " + e.getMessage(), e);
        }
    }

    /**
     * 解析 Word 文档（docx）
     */
    private List<Document> parseWord(String docUrl, Consumer<ParseProgress> progressCallback) {
        try {
            progressCallback.accept(new ParseProgress(0, 0, "", "正在解析 Word 文档...", 30));
            
            List<Document> documents = new ArrayList<>();
            try (InputStream inputStream = new URL(docUrl).openStream();
                 XWPFDocument document = new XWPFDocument(inputStream)) {
                
                StringBuilder content = new StringBuilder();
                for (XWPFParagraph paragraph : document.getParagraphs()) {
                    String text = paragraph.getText();
                    if (StringUtils.hasText(text)) {
                        content.append(text).append("\n");
                    }
                }
                
                if (content.length() > 0) {
                    // 将整个 Word 文档作为一个 Document，或者按段落分割
                    // 这里按段落分割，每个段落作为一个 Document
                    String[] paragraphs = content.toString().split("\n\n");
                    for (int i = 0; i < paragraphs.length; i++) {
                        String para = paragraphs[i].trim();
                        if (StringUtils.hasText(para)) {
                            Map<String, Object> metadata = new HashMap<>();
                            metadata.put("paragraphIndex", i);
                            documents.add(new Document(para, metadata));
                        }
                    }
                }
            }
            
            progressCallback.accept(new ParseProgress(0, 0, "", "Word 解析完成", 50));
            return documents;
        } catch (Exception e) {
            log.error("Word 解析失败，docUrl: {}", docUrl, e);
            throw new RuntimeException("Word 解析失败: " + e.getMessage(), e);
        }
    }

    /**
     * 从 URL 获取文件类型
     */
    private String getFileTypeFromUrl(String url) {
        if (!StringUtils.hasText(url)) {
            return "unknown";
        }
        String lowerUrl = url.toLowerCase();
        if (lowerUrl.endsWith(".pdf")) {
            return "pdf";
        } else if (lowerUrl.endsWith(".md") || lowerUrl.endsWith(".markdown")) {
            return "md";
        } else if (lowerUrl.endsWith(".docx")) {
            return "docx";
        } else if (lowerUrl.endsWith(".doc")) {
            return "doc";
        }
        return "unknown";
    }
}

